import pandas as pd
import os
import math

'''
    获取存储的day、flow、涨停率
'''
def get_origin_csv_file():
    max_w_df = pd.read_csv('input/week_max_40.csv',index_col=0)[['code','max_close',]]
    days_df = pd.read_csv('input/day.csv',index_col=0)
    flow_df = pd.read_csv('input/flow.csv',index_col=0)
    zt_ratio_df = pd.read_csv('input/flow.csv',index_col=0)
    return days_df,flow_df,zt_ratio_df,max_w_df


def get_stock_list():
    return pd.read_csv('input/stock_list.csv',index_col=0)



"""
    将CSV文件拆分为指定数量的小文件
    :param input_file: 输入CSV文件路径
    :param output_dir: 输出目录路径
    :param num_files: 目标拆分文件数量
"""
def split_csv(input_file, output_dir, num_files):
   
    # 检查输出目录是否存在，不存在则创建
    os.makedirs(output_dir, exist_ok=True)
    
    # 读取文件总行数（不加载全部数据）
    total_rows = sum(1 for _ in open(input_file, 'r')) - 1  # 减去表头行
    rows_per_file = math.ceil(total_rows / num_files)  # 计算每个文件行数
    
    # 分块读取并写入文件
    for i, chunk in enumerate(pd.read_csv(input_file, chunksize=rows_per_file)):
        output_path = os.path.join(output_dir, f"part_{i+1}.csv")
        print(output_path)
        chunk.to_csv(output_path, index=False)
        print(f"已生成: {output_path} ({len(chunk)}行)")



"""
    合并目录下所有CSV文件
    :param input_dir: 输入目录路径（包含多个CSV）
    :param output_file: 合并后的输出文件路径
"""
def merge_csv(input_dir, output_file):
   
    all_files = [
        os.path.join(input_dir, f) 
        for f in os.listdir(input_dir) 
        if f.endswith(".csv")
    ]
    
    # 读取并合并文件
    df_list = [pd.read_csv(file) for file in all_files]
    merged_df = pd.concat(df_list, ignore_index=True)
    
    # 保存合并结果
    merged_df.to_csv(output_file, index=False)
    print(f"已合并 {len(all_files)} 个文件 → {output_file}")




# # 拆分文件（分成3个小文件）
# split_csv("large_data.csv", "./split_files", num_files=3)

# # 合并文件
# merge_csv("./split_files", "combined_data.csv")



import pandas as pd

def split_and_save(df: pd.DataFrame, 
                  output_dir: str = "output",
                  prefix: str = "part") -> None:
    """
    将DataFrame等分3份存储到指定目录
    
    参数:
        df: 要分割的DataFrame
        output_dir: 输出目录路径（默认output）
        prefix: 文件名前缀（默认part）
    
    特性:
        1. 自动处理余数分配[1](@ref)
        2. 直接使用字符串路径拼接
        3. 依赖目录已存在（需手动创建output_dir）
    """
    total = len(df)
    base, rem = divmod(total, 3)
    
    # 生成分片索引（网页1的iloc分割法）
    slices = [
        (0, base + rem),
        (base + rem, base*2 + rem),
        (base*2 + rem, total)
    ]
    
    # 直接存储CSV文件（网页2的批量存储逻辑）
    for i, (s, e) in enumerate(slices, 1):
        df.iloc[s:e].to_csv(
            f"{output_dir}/{prefix}_{i}.csv", 
            index=False
        )

def load_and_merge(input_dir: str = "input",
                  prefix: str = "part") -> pd.DataFrame:
    """
    合并指定目录的分片文件
    
    参数:
        input_dir: 输入目录路径（默认input）
        prefix: 文件名前缀（默认part）
    
    返回:
        合并后的完整DataFrame
    
    特性:
        1. 预设固定分片数[2](@ref)
        2. 直接构造文件路径
    """
    # 生成固定文件名列表
    files = [f"{input_dir}/{prefix}_{i}.csv" for i in [1,2,3]]
    
    # 合并数据（网页2的concat方法）
    return pd.concat(
        (pd.read_csv(f) for f in files),
        ignore_index=True
    )






'''  切分文件,切分Dataframe对象 '''
def split_dataframe(df, num_files):
    chunks = np.array_split(df, num_files)
    for i, chunk in enumerate(chunks):
        chunk.to_csv('input/part2_'+str(i+1)+'.csv', index=False)




